"""
pip install llama_index
pip install llama-index-embeddings-huggingface
pip install llama-index-llms-langchain

"""
import os
import shutil

import chromadb
from dotenv import load_dotenv
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from langchain_openai import ChatOpenAI
from llama_index.core import Settings, StorageContext, SimpleDirectoryReader, VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore

# 自动加载 .env 文件中的环境变量
load_dotenv()

# ===============================
# 2️⃣ 配置向量模型（Embedding）
# ===============================
Settings.embed_model = HuggingFaceEmbedding(
    model_name="D:/models/BAAI_bge-small-zh-v1.5"
)

# ===============================
# 3️⃣ 配置大语言模型（LLM）
# ===============================
Settings.llm = ChatOpenAI(api_key=os.getenv("XUNFEI_API_KEY"), openai_api_base=os.getenv("XUNFEI_API_BASE"), model="xdeepseekv32exp", streaming=True, temperature=0.7)

# ===============================
# 4️⃣ 初始化 Chroma 向量数据库
# ===============================
# 文本数据目录（里面放要索引的文档）
data_dir = "data"
# 向量数据库存储目录
persist_dir = "./chroma_db"

## 测试使用因为要多次运行，每次都会重新创建索引，所以如果检测到 data/ 下的文件有变化，就删除旧索引目录
if os.path.exists(persist_dir):
    shutil.rmtree(persist_dir)

# 创建一个持久化的 Chroma 客户端，存储在指定路径
db = chromadb.PersistentClient(path=persist_dir)
# 获取或创建一个集合（相当于一个向量表）
chroma_collection = db.get_or_create_collection("quick_start")
# 将该集合封装成 llama_index 可用的向量存储对象
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
# 创建存储上下文，用于保存索引的元信息和向量数据
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# ===============================
# 5️⃣ 加载本地文档
# ===============================
# SimpleDirectoryReader 会自动读取指定文件夹下的所有文本文件（.txt, .md, .pdf等）
# 加载 data 目录中的所有文档，返回文档对象列表
documents = SimpleDirectoryReader(data_dir).load_data()

# ===============================
# 6️⃣ 创建向量索引
# ===============================
# 将文档向量化并存入向量数据库中
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context
)


index.storage_context.persist()


# ===============================
# 7️⃣ 创建查询引擎并提问
# ===============================
query_engine = index.as_query_engine()  # 将索引转为查询引擎，用于自然语言问答
response = query_engine.query("王德宝是谁")  # 询问一个问题，底层会用向量检索 + LLM 总结回答
print(f"回答: {response}")
